package com.geccocrawler.socks5.handler.ss5;

import com.geccocrawler.socks5.handler.ss5.forward.Common;
import com.geccocrawler.socks5.handler.ss5.forward.tcp.TcpClient2DestHandler;
import com.geccocrawler.socks5.handler.ss5.forward.tcp.TcpDest2ClientHandler;
import com.geccocrawler.socks5.handler.ss5.forward.udp.UdpDest2ClientHandler;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioDatagramChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.socksx.v5.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetSocketAddress;

public class Socks5CommandRequestHandler extends SimpleChannelInboundHandler<DefaultSocks5CommandRequest> {
    static EventLoopGroup bossGroup = new NioEventLoopGroup();

    private int udpPort;

    public Socks5CommandRequestHandler(int udpPort) {
        this.udpPort = udpPort;
    }

    private static final Logger logger = LoggerFactory.getLogger(Socks5CommandRequestHandler.class);

    @Override
    protected void channelRead0(final ChannelHandlerContext clientChannelContext, DefaultSocks5CommandRequest msg) throws Exception {
        logger.debug("目标服务器  : " + msg.type() + "," + msg.dstAddr() + "," + msg.dstPort());
        //tcp
        if (msg.type() == Socks5CommandType.CONNECT) {
            tcpProxyRequestHandler(clientChannelContext, msg);
        } else if (msg.type() == Socks5CommandType.UDP_ASSOCIATE) {
            udpProxyRequestHandler(clientChannelContext, msg);
        } else {
            clientChannelContext.fireChannelRead(msg);
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
//        logger.info("通道已关闭! ctx: {}", ctx);
    }

    private void tcpProxyRequestHandler(final ChannelHandlerContext clientChannelContext, DefaultSocks5CommandRequest msg) throws InterruptedException {
        logger.trace("TCP->准备连接目标服务器");
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(bossGroup)
                .channel(NioSocketChannel.class)
                .option(ChannelOption.TCP_NODELAY, true)
                .handler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
//                            ch.pipeline().addLast("timeout", new IdleStateHandler(8, 8, 0, TimeUnit.SECONDS) {
//                                @Override
//                                public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
//                                    if(IdleStateEvent.class.isAssignableFrom(evt.getClass())){
//                                        IdleStateEvent event = (IdleStateEvent) evt;
//                                    }
//                                }
//                            });
                        ch.pipeline().addLast(new TcpDest2ClientHandler(clientChannelContext));
                    }

                    @Override
                    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
                        logger.info("代理通道已关闭! ctx: {}", ctx);
                    }
                });

        ChannelFuture future = bootstrap.connect(msg.dstAddr(), msg.dstPort()).sync();
        future.addListener((ChannelFutureListener) future1 -> {
            if (future1.isSuccess()) {
                logger.trace("成功连接目标服务器");
                clientChannelContext.pipeline().addLast(new TcpClient2DestHandler(future1));
                Socks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, Socks5AddressType.IPv4);
                clientChannelContext.writeAndFlush(commandResponse);
            } else {
                Socks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.FAILURE, Socks5AddressType.IPv4);
                clientChannelContext.writeAndFlush(commandResponse);
            }
        });
    }

    private void udpProxyRequestHandler(final ChannelHandlerContext clientChannelContext, DefaultSocks5CommandRequest msg) throws InterruptedException {
        logger.trace("Udp->准备连接目标服务器");
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(bossGroup).channel(NioDatagramChannel.class)
//                .option(ChannelOption.SO_RCVBUF, 64 * 1024)// 设置UDP读缓冲区为64k
//                .option(ChannelOption.SO_SNDBUF, 64 * 1024)// 设置UDP写缓冲区为64k
                .handler(new ChannelInitializer<NioDatagramChannel>() {
                    @Override
                    protected void initChannel(NioDatagramChannel ch) throws Exception {
                        ch.pipeline().addLast(new UdpDest2ClientHandler(clientChannelContext));
                    }
                });

        ChannelFuture channelFuture = bootstrap.bind(0).sync();
        channelFuture.addListener((ChannelFutureListener) future -> {
            if (future.isSuccess()) {
                future.channel().attr(Common.REMOTE_DES).set(new InetSocketAddress(msg.dstAddr(), msg.dstPort()));
                Socks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(
                        Socks5CommandStatus.SUCCESS,
                        Socks5AddressType.IPv4,
                        ((InetSocketAddress)clientChannelContext.channel().localAddress()).getAddress().getHostAddress(),
                        udpPort);
                clientChannelContext.writeAndFlush(commandResponse);
            } else {
                Socks5CommandResponse commandResponse = new DefaultSocks5CommandResponse(Socks5CommandStatus.FAILURE, Socks5AddressType.IPv4);
                clientChannelContext.writeAndFlush(commandResponse);
            }
        });
    }
}
